#include "Astar_searcher.h"
#include <math.h>

using namespace std;
using namespace Eigen;

void AstarPathFinder::initGridMap(double _resolution, Vector3d global_xyz_l, Vector3d global_xyz_u, int max_x_id, int max_y_id, int max_z_id)
{   
    gl_xl = global_xyz_l(0);
    gl_yl = global_xyz_l(1);
    gl_zl = global_xyz_l(2);

    gl_xu = global_xyz_u(0);
    gl_yu = global_xyz_u(1);
    gl_zu = global_xyz_u(2);
    
    GLX_SIZE = max_x_id;
    GLY_SIZE = max_y_id;
    GLZ_SIZE = max_z_id;
    GLYZ_SIZE  = GLY_SIZE * GLZ_SIZE;
    GLXYZ_SIZE = GLX_SIZE * GLYZ_SIZE;

    resolution = _resolution;
    inv_resolution = 1.0 / _resolution;    

    data = new uint8_t[GLXYZ_SIZE];
    memset(data, 0, GLXYZ_SIZE * sizeof(uint8_t));
    
    GridNodeMap = new GridNodePtr ** [GLX_SIZE];
    for(int i = 0; i < GLX_SIZE; i++){
        GridNodeMap[i] = new GridNodePtr * [GLY_SIZE];
        for(int j = 0; j < GLY_SIZE; j++){
            GridNodeMap[i][j] = new GridNodePtr [GLZ_SIZE];
            for( int k = 0; k < GLZ_SIZE;k++){
                Vector3i tmpIdx(i,j,k);
                Vector3d pos = gridIndex2coord(tmpIdx);
                GridNodeMap[i][j][k] = new GridNode(tmpIdx, pos);
            }
        }
    }
}

void AstarPathFinder::resetGrid(GridNodePtr ptr)
{
    ptr->id = 0;
    ptr->cameFrom = NULL;
    ptr->gScore = inf;
    ptr->fScore = inf;
}

void AstarPathFinder::resetUsedGrids()
{   
    for(int i=0; i < GLX_SIZE ; i++)
        for(int j=0; j < GLY_SIZE ; j++)
            for(int k=0; k < GLZ_SIZE ; k++)
                resetGrid(GridNodeMap[i][j][k]);
}

void AstarPathFinder::setObs(const double coord_x, const double coord_y, const double coord_z)
{
    if( coord_x < gl_xl  || coord_y < gl_yl  || coord_z <  gl_zl || 
        coord_x >= gl_xu || coord_y >= gl_yu || coord_z >= gl_zu )
        return;

    int idx_x = static_cast<int>( (coord_x - gl_xl) * inv_resolution);
    int idx_y = static_cast<int>( (coord_y - gl_yl) * inv_resolution);
    int idx_z = static_cast<int>( (coord_z - gl_zl) * inv_resolution);      

    data[idx_x * GLYZ_SIZE + idx_y * GLZ_SIZE + idx_z] = 1;
}

vector<Vector3d> AstarPathFinder::getVisitedNodes()
{   
    vector<Vector3d> visited_nodes;
    for(int i = 0; i < GLX_SIZE; i++)
        for(int j = 0; j < GLY_SIZE; j++)
            for(int k = 0; k < GLZ_SIZE; k++){   
                //if(GridNodeMap[i][j][k]->id != 0) // visualize all nodes in open and close list
                if(GridNodeMap[i][j][k]->id == -1)  // visualize nodes in close list only
                    visited_nodes.push_back(GridNodeMap[i][j][k]->coord);
            }

    ROS_WARN("visited_nodes size : %d", visited_nodes.size());
    return visited_nodes;
}

Vector3d AstarPathFinder::gridIndex2coord(const Vector3i & index) 
{
    Vector3d pt;

    pt(0) = ((double)index(0) + 0.5) * resolution + gl_xl;
    pt(1) = ((double)index(1) + 0.5) * resolution + gl_yl;
    pt(2) = ((double)index(2) + 0.5) * resolution + gl_zl;

    return pt;
}

Vector3i AstarPathFinder::coord2gridIndex(const Vector3d & pt) 
{
    Vector3i idx;
    idx <<  min( max( int( (pt(0) - gl_xl) * inv_resolution), 0), GLX_SIZE - 1),
            min( max( int( (pt(1) - gl_yl) * inv_resolution), 0), GLY_SIZE - 1),
            min( max( int( (pt(2) - gl_zl) * inv_resolution), 0), GLZ_SIZE - 1);                  
  
    return idx;
}

Eigen::Vector3d AstarPathFinder::coordRounding(const Eigen::Vector3d & coord)
{
    return gridIndex2coord(coord2gridIndex(coord));
}
/**
 * @brief 判断参数指定的索引对应的节点是否是障碍物
 * 
 * @param index 
 * @return true 
 * @return false 
 */
inline bool AstarPathFinder::isOccupied(const Eigen::Vector3i & index) const
{
    return isOccupied(index(0), index(1), index(2));
}
/**
 * @brief 判断参数指定的索引对应的节点是否空白
 * 
 * @param index 
 * @return true 
 * @return false 
 */
inline bool AstarPathFinder::isFree(const Eigen::Vector3i & index) const
{
    return isFree(index(0), index(1), index(2));
}
/**
 * @brief 判断参数指定的索引对应的节点是否是障碍物
 * 
 * @param idx_x 
 * @param idx_y 
 * @param idx_z 
 * @return true 
 * @return false 
 */
inline bool AstarPathFinder::isOccupied(const int & idx_x, const int & idx_y, const int & idx_z) const 
{
    return  (idx_x >= 0 && idx_x < GLX_SIZE && idx_y >= 0 && idx_y < GLY_SIZE && idx_z >= 0 && idx_z < GLZ_SIZE && 
            (data[idx_x * GLYZ_SIZE + idx_y * GLZ_SIZE + idx_z] == 1));
}
// 判断参数指定的索引对应的节点是否空白
inline bool AstarPathFinder::isFree(const int & idx_x, const int & idx_y, const int & idx_z) const 
{
    return (idx_x >= 0 && idx_x < GLX_SIZE && idx_y >= 0 && idx_y < GLY_SIZE && idx_z >= 0 && idx_z < GLZ_SIZE && 
           (data[idx_x * GLYZ_SIZE + idx_y * GLZ_SIZE + idx_z] < 1));
}
/**
 * @brief 展开当前节点，获取可访问的邻居节点和当前节点与邻居节点之间的距离
 * 
 * @param currentPtr 当前节点
 * @param neighborPtrSets 可访问的邻居节点数组
 * @param edgeCostSets 可访问的邻居节点与当前节点之间距离的数组
 */
inline void AstarPathFinder::AstarGetSucc(GridNodePtr currentPtr, vector<GridNodePtr> & neighborPtrSets, vector<double> & edgeCostSets)
{   
    neighborPtrSets.clear();
    edgeCostSets.clear();
    /*
    *
    STEP 4: finish AstarPathFinder::AstarGetSucc yourself 
    please write your code below
    *
    *
    */
     if(currentPtr == nullptr)
        std::cout << "Error: Current pointer is null!" << endl;


    int this_x = currentPtr->index[0];
    int this_y = currentPtr->index[1];
    int this_z = currentPtr->index[2];
    auto this_coord = currentPtr->coord;

    // 当前要访问的节点在地图上的索引
    int  n_x, n_y, n_z;
    double dist;
    // 要处理的邻居节点
    GridNodePtr temp_ptr = nullptr;
    // 邻居节点的坐标
    Eigen::Vector3d n_coord;

    // 遍历周围点,获取周围点的edgeCostSets
    for(int i = -1; i <= 1; ++i ){
        for(int j = -1; j <= 1; ++j ){
            for(int k = -1; k <= 1; ++k){
                // 如果是当前点，跳过
                if( i == 0 && j == 0 && k == 0)
                    continue;

                // 获取要访问的节点的索引
                n_x = this_x + i;
                n_y = this_y + j;
                n_z = this_z + k;

                // 如果索引越界
                if( (n_x < 0) || (n_x > (GLX_SIZE - 1)) || (n_y < 0) || (n_y > (GLY_SIZE - 1) ) || (n_z < 0) || (n_z > (GLZ_SIZE - 1)))
                    continue;

                // 如果是障碍物，跳过
                if(isOccupied(n_x, n_y, n_z))
                    continue;

                // 获取一个邻居节点
                temp_ptr = GridNodeMap[n_x][n_y][n_z];

                // 节点已被扩展
                if(temp_ptr->id == -1) 
                    continue;

                n_coord = temp_ptr->coord;

                if(temp_ptr == currentPtr){
                    std::cout << "Error: temp_ptr == currentPtr)" << std::endl;
                }

                // if( (std::abs(n_coord[0] - this_coord[0]) < 1e-6) and (std::abs(n_coord[1] - this_coord[1]) < 1e-6) and (std::abs(n_coord[2] - this_coord[2]) < 1e-6 )){
                //     std::cout << "Error: Not expanding correctly!" << std::endl;
                //     std::cout << "n_coord:" << n_coord[0] << " "<<n_coord[1]<<" "<<n_coord[2] << std::endl;
                //     std::cout << "this_coord:" << this_coord[0] << " "<<this_coord[1]<<" "<<this_coord[2] << std::endl;

                //     std::cout << "current node index:" << this_x << " "<< this_y<<" "<< this_z << std::endl;
                //     std::cout << "neighbor node index:" << n_x << " "<< n_y<<" "<< n_z << std::endl;
                // }

                dist = std::sqrt( (n_coord[0] - this_coord[0]) * (n_coord[0] - this_coord[0])+
                        (n_coord[1] - this_coord[1]) * (n_coord[1] - this_coord[1])+
                        (n_coord[2] - this_coord[2]) * (n_coord[2] - this_coord[2]));
                
                neighborPtrSets.push_back(temp_ptr); 
                edgeCostSets.push_back(dist);
            }
        }
    }
}

double AstarPathFinder::getHeu(GridNodePtr node1, GridNodePtr node2)
{
    /* 
    choose possible heuristic function you want
    Manhattan, Euclidean, Diagonal, or 0 (Dijkstra)
    Remember tie_breaker learned in lecture, add it here ?
    *
    *
    *
    STEP 1: finish the AstarPathFinder::getHeu , which is the heuristic function
    please write your code below
    *
    *
    */

    //  Euclidean Distance
    double dis = 0;
    if(heu_type == 0) {
        dis = sqrt(pow(node1->coord.x() - node2->coord.x(),2) + pow(node1->coord.y() - node2->coord.y(), 2)
                 + pow(node1->coord.z() - node2->coord.z(), 2));
    } else if(heu_type == 1) {
    // 曼哈顿距离算法

        dis = abs(node1->coord.x() - node2->coord.x()) + abs(node1->coord.y() - node2->coord.y())
            + abs(node1->coord.z() - node2->coord.z());
    }
    // ROS_INFO("AstarPathFinder::getHeu---->dis %f", dis);
    return dis;
}

void AstarPathFinder::AstarGraphSearch(Vector3d start_pt, Vector3d end_pt)
{   
    ros::Time time_1 = ros::Time::now();    

    //index of start_point and end_point
    Vector3i start_idx = coord2gridIndex(start_pt);
    Vector3i end_idx   = coord2gridIndex(end_pt);
    goalIdx = end_idx;

    //position of start_point and end_point
    start_pt = gridIndex2coord(start_idx);
    end_pt   = gridIndex2coord(end_idx);

    //Initialize the pointers of struct GridNode which represent start node and goal node
    GridNodePtr startPtr = new GridNode(start_idx, start_pt);
    GridNodePtr endPtr   = new GridNode(end_idx,   end_pt);

    //openSet is the open_list implemented through multimap in STL library
    openSet.clear();
    // currentPtr represents the node with lowest f(n) in the open_list
    GridNodePtr currentPtr  = NULL;
    GridNodePtr neighborPtr = NULL;

    // cout << "111111111111111111111111111" << endl;
    //put start node in open set
    startPtr -> gScore = 0;
    startPtr -> fScore = getHeu(startPtr,endPtr);   
    //STEP 1: finish the AstarPathFinder::getHeu , which is the heuristic function
    startPtr -> id = 1; 
    startPtr -> coord = start_pt;
    openSet.insert( make_pair(startPtr -> fScore, startPtr) );

    // cout << "222222222222222222222222222222" << endl;
    /*
    *
    STEP 2 :  some else preparatory works which should be done before while loop
    please write your code below
    *
    *
    */
    // 邻居节点数组
    vector<GridNodePtr> neighborPtrSets;
    // 当前节点到邻居节点之间的代价值数组
    vector<double> edgeCostSets;

    // this is the main loop
    while ( !openSet.empty() ){
        /*
         *
         *
         step 3: Remove the node with lowest cost function from open set to closed set
         please write your code below
        
         IMPORTANT NOTE!!!
         This part you should use the C++ STL: multimap, more details can be find in Homework description
         *
         *
         */
        // openset:待访问节点容器;closed set:访问过节点容器

        // cout << "step 3------->11111111111" << endl;
        int x = openSet.begin()->second->index.x(); 
        int y = openSet.begin()->second->index.y(); 
        int z = openSet.begin()->second->index.z();
        openSet.erase(openSet.begin());
        currentPtr = GridNodeMap[x][y][z];

        // 如果节点被访问过;则返回
        if(currentPtr->id == -1)
            continue;
        // 标记id为-1,表示节点已被扩展
        currentPtr->id = -1;
        // cout << "step 3------->2222222222222" << endl;

        // if the current node is the goal 
        if( currentPtr->index == goalIdx ){
            ros::Time time_2 = ros::Time::now();
            terminatePtr = currentPtr;
            ROS_WARN("[A*]{sucess}  Time in A*  is %f ms, path cost if %f m", (time_2 - time_1).toSec() * 1000.0, currentPtr->gScore * resolution );            
            return;
        }

        //get the succetion
        // cout << "step 4------->11111111111111-->AstarGetSucc" << endl;
        AstarGetSucc(currentPtr, neighborPtrSets, edgeCostSets);  
        //STEP 4: finish AstarPathFinder::AstarGetSucc yourself     

        // cout << "step 4------->2222222222222-->AstarGetSucc" << endl;
        /*
        *
        *
        STEP 5:  For all unexpanded neigbors "m" of node "n", please finish this for loop
        please write your code below
        *        
        */         
        for(int i = 0; i < (int)neighborPtrSets.size(); i++){
            /*
             *
             *
             Judge if the neigbors have been expanded
             please write your code below
            
             IMPORTANT NOTE!!!
             neighborPtrSets[i]->id = -1 : expanded, equal to this node is in close set
             neighborPtrSets[i]->id = 1 : unexpanded, equal to this node is in open set
             *        
             */
            neighborPtr = neighborPtrSets[i];
            if(neighborPtr -> id == 0){ //discover a new node, which is not in the closed set and open set
                /*
                 *
                 *
                 STEP 6:  As for a new node, do what you need do ,and then put neighbor in open set and record it
                 please write your code below
                 *        
                 */
                // 设置节点的g(n)值和f(n)值
                // cout << "step 5------->11111111111111" << endl;
                neighborPtr->gScore = currentPtr->gScore + edgeCostSets[i];
                neighborPtr->fScore = neighborPtr->gScore + getHeu(neighborPtr,endPtr);
                neighborPtr->cameFrom = currentPtr;
                // push node "m" into OPEN
                openSet.insert(make_pair(neighborPtr -> fScore, neighborPtr));
                // 把节点设置为未扩展
                neighborPtr -> id = 1;
                // cout << "step 5------->2222222222222222" << endl;

                continue;
            }
            else if(neighborPtr -> id == 1){ //this node is in open set and need to judge if it needs to update, the "0" should be deleted when you are coding
                /*
                 *
                 *
                 STEP 7:  As for a node in open set, update it , maintain the openset ,and then put neighbor in open set and record it
                 please write your code below
                 *        
                 */
                // cout << "step 7------->11111111111111" << endl;
                // 如果邻居节点 m 的代价值 大于 当前节点 n 的代价 + Cnm，更新邻居节点代价值
                if( neighborPtr->gScore > (currentPtr->gScore + edgeCostSets[i]))
                {
                    neighborPtr -> gScore = currentPtr -> gScore + edgeCostSets[i];
                    neighborPtr -> fScore = neighborPtr -> gScore + getHeu(neighborPtr,endPtr);
                    neighborPtr -> cameFrom = currentPtr;
                }
               
                // cout << "step 7------->2222222222222222" << endl;
                continue;
            }
            else{//this node is in closed set
                /*
                 *
                 please write your code below
                 *        
                 */
                // 什么也不做
                continue;
            }
        }      
    }
    //if search fails
    ros::Time time_2 = ros::Time::now();
    if((time_2 - time_1).toSec() > 0.1)
        ROS_WARN("Time consume in Astar path finding is %f", (time_2 - time_1).toSec() );
}


vector<Vector3d> AstarPathFinder::getPath() 
{   
    vector<Vector3d> path;
    vector<GridNodePtr> gridPath;
    /*
     *
     *
     STEP 8:  trace back from the curretnt nodePtr to get all nodes along the path
     please write your code below
     *      
     */

    // cout << "step 8------->11111111111111------------>getPath" << endl;
    auto ptr = terminatePtr;
    while(ptr -> cameFrom != NULL)
    {
        gridPath.push_back(ptr);
        ptr = ptr->cameFrom;
    }

    // cout << "step 8------->22222222222222------------>getPath" << endl;

    for (auto ptr: gridPath)
        path.push_back(ptr->coord);
        
    reverse(path.begin(),path.end());

    return path;
}


// if the difference of f is trivial, then choose then prefer the path along the straight line from start to goal
// discared!!! 
GridNodePtr & AstarPathFinder::TieBreaker(const std::multimap<double, GridNodePtr> &  openSet, const GridNodePtr & endPtr)
{
    // todo do I have to update the f in openSet??
    std::multimap<double, GridNodePtr> local_set;

    auto f_min = openSet.begin()->first;
    auto f_max = f_min + 1e-2;
    auto itlow = openSet.lower_bound (f_min);
    auto itup = openSet.upper_bound(f_max);
    double cross, f_new;

    for (auto it=itlow; it!=itup; ++it)
    {
        std::cout << "f value is:" << (*it).first << " pointer is: " << (*it).second << '\n';
        cross = std::abs(endPtr->coord(0) - (*it).second->coord(0)) +
                std::abs(endPtr->coord(1) - (*it).second->coord(1)) +
                std::abs(endPtr->coord(2) - (*it).second->coord(2));
        f_new = (*it).second->fScore + 0.001 * cross;
        local_set.insert( make_pair(f_new, (*it).second) ); // todo what is iterator, is this way correct?
    }


    return local_set.begin()->second;
}